from math import cos, hypot, pi, sin
import new
from random import Random, randint, uniform

import pyglet
from pyglet.gl import *

import src.shader

def _buildBallShader():
	ball_shader = src.shader.ShaderProgram(
	 src.shader.VertexShader("""
		void main(){
			gl_Position = gl_ModelViewProjectionMatrix * gl_Vertex;
		}"""
	 ),
	 src.shader.FragmentShader("""
		uniform vec3 rgb;
		
		void main(){
			gl_FragColor = vec4(rgb[0], rgb[1], rgb[2], 1.0);
		}"""
	 )
	)
	_shader_message = ball_shader.link()
	if _shader_message:
		print _shader_message
	return ball_shader
	
def _buildTrailShader():
	trail_shader = src.shader.ShaderProgram(
	 src.shader.VertexShader("""
		void main(){
			gl_Position = gl_ModelViewProjectionMatrix * gl_Vertex;
		}"""
	 ),
	 src.shader.FragmentShader("""
		uniform vec3 rgb;
		uniform float a;
		
		void main(){
			gl_FragColor = vec4(rgb[0], rgb[1], rgb[2], a);
		}"""
	 )
	)
	_shader_message = trail_shader.link()
	if _shader_message:
		print _shader_message
	return trail_shader
	
def _buildBallVertices(vertex_count, radius):
	vertices = [0.0, 0.0]
	for i in range(0, 360, 360 / vertex_count):
		i_rad = i * (pi / 180.0)
		vertices += [radius * cos(i_rad), radius * sin(i_rad)]
	vertices += vertices[2:4] #Close the shape.
	return (GLfloat * len(vertices))(*vertices)
	
	
class EnergyBall(object):
	x = None
	y = None
	vec_x = None
	vec_y = None
	
	hidden = False
	fixed_velocity = False
	
	def __init__(self):
		self.x = 0.0
		self.y = 0.0
		self.vec_x = 0.0
		self.vec_y = 0.0
		self.charge = 0.0
		
	def push(self, vec_x, vec_y, pulled=False):
		self.vec_x = vec_x
		self.vec_y = vec_y
		self.fixed_velocity = pulled
		
	def tick(self, width, height, rnd):
		if not self.hidden:
			old_x = self.x
			old_y = self.y
			
			self.x += self.vec_x
			if self.x <= 0.0:
				self.x = 0.0
				self.vec_x *= -1.0
			elif self.x >= width:
				self.x = width
				self.vec_x *= -1.0
			self.y += self.vec_y
			if self.y <= 0.0:
				self.y = 0.0
				self.vec_y *= -1.0
			elif self.y >= height:
				self.y = height
				self.vec_y *= -1.0
				
			if not self.fixed_velocity:
				if abs(self.vec_x) > self.target_speed_x:
					self.vec_x *= min(1.0, self.deceleration_buffer + (1.0 - self.charge / self.max_charge))
				else:
					self.vec_x += rnd.uniform(-self.speed_variance_x, self.speed_variance_x)
				if abs(self.vec_y) > self.target_speed_y:
					self.vec_y *= min(1.0, self.deceleration_buffer + (1.0 - self.charge / self.max_charge))
				else:
					self.vec_y += rnd.uniform(-self.speed_variance_y, self.speed_variance_y)
					
			self.charge = min(self.max_charge, self.charge + self.charge_rate)
			
			return [old_x, old_y, self.x, self.y]
		return None
		
	def drain(self, amount):
		self.charge = max(0.0, self.charge - amount)
		return not self.charge == 0.0
		
	def draw(self):
		if not self.hidden:
			glLoadIdentity()
			glTranslatef(self.x, self.y, 0.0)
			glDrawArrays(GL_TRIANGLE_FAN, 0, self.vertex_count)
			
	def hide(self):
		self.hidden = True
		
	def show(self):
		self.hidden = False
		
	def release(self):
		self.fixed_velocity = False
		
class World(object):
	_energy_balls = None
	_energy_balls_types = None
	_energy_balls_by_type = None
	_energy_balls_trails_by_type = None
	
	_rnd = None
	
	_target_speed = None
	_target_speed_variance = None
	_deceleration_buffer = None
	_trail_length = None
	
	_width = None
	_height = None
	
	_ball_vertices = None
	_ball_shader = None
	_trail_shader = None
	
	def __init__(self, width, height, seed, ball_types):
		assert type(width) == int
		assert type(height) == int
		assert type(ball_types) in (tuple, list)
		
		self._width = width
		self._height = height
		self._rnd = Random(seed)
		
		self._energy_balls = []
		self._energy_balls_types = {}
		self._energy_balls_by_type = {}
		self._energy_balls_trails_by_type = {}
		
		self._ball_shader = _buildBallShader()
		self._trail_shader = _buildTrailShader()
		
		for ball_type in ball_types:
			assert len(ball_type) == 16
			(name, ball_class, count,
			 vertex_count, radius,
			 colour, trail_colour, trail_length, trail_brightness,
			 target_speed_x, target_speed_y, speed_variance_x, speed_variance_y, deceleration_buffer,
			 max_charge, charge_rate
			) = ball_type
			
			assert isinstance(name, basestring)
			assert type(count) == int
			assert type(vertex_count) == int
			assert type(radius) in (float, int)
			assert type(colour) in (tuple, list) and len(colour) == 3 and not [None for c in colour if not type(c) in(float, int)]
			assert type(trail_colour) in (tuple, list) and len(trail_colour) == 3 and not [None for c in trail_colour if not type(c) in (float, int)]
			assert type(trail_length) == int
			assert type(trail_brightness) in (float, int)
			assert type(target_speed_x) in (float, int)
			assert type(target_speed_y) in (float, int)
			assert type(speed_variance_x) in (float, int)
			assert type(speed_variance_y) in (float, int)
			assert type(deceleration_buffer) in (float, int)
			assert type(max_charge) in (float, int)
			assert type(charge_rate) in (float, int)
			
			ball_vertices = _buildBallVertices(max(3, int(vertex_count)), max(0.5, float(radius)))
			ball_type = new.classobj(name, (ball_class,), {
			 'ball_vertices': ball_vertices,
			 'vertex_count': len(ball_vertices) // 2,
			 'colour': tuple([float(c) for c in colour]),
			 'trail_colour': tuple([float(c) for c in trail_colour]),
			 'trail_length': max(0, int(trail_length)),
			 'trail_brightness': max(0.0, min(1.0, float(trail_brightness))),
			 'target_speed_x': max(0.0, float(target_speed_x)),
			 'target_speed_y': max(0.0, float(target_speed_y)),
			 'speed_variance_x': max(0.0, float(speed_variance_x)),
			 'speed_variance_y': max(0.0, float(speed_variance_y)),
			 'deceleration_buffer': max(0.0, min(1.0, float(deceleration_buffer))),
			 'max_charge': max(0.0, float(max_charge)),
			 'charge_rate': max(0.0, float(charge_rate)),
			})
			self._energy_balls_types[name] = ball_type
			
			energy_balls = []
			for i in xrange(count):
				energy_ball = ball_type()
				self._energy_balls.append(energy_ball)
				energy_balls.append(energy_ball)
			self._energy_balls_by_type[ball_type] = tuple(energy_balls)
			self._energy_balls_trails_by_type[ball_type] = []
		self._rnd.shuffle(self._energy_balls)
		self._energy_balls = tuple(self._energy_balls)
		
	def randomizeBallLocations(self, min_x=0.0, max_x=None, min_y=0.0, max_y=None, ball_type=None):
		assert type(min_x) in (int, float)
		assert max_x is None or type(max_x) in (int, float)
		assert type(min_y) in (int, float)
		assert max_y is None or type(max_y) in (int, float)
		assert ball_type is None or isinstance(ball_type, basestring)
		
		if max_x is None:
			max_x = float(self._width)
		if max_y is None:
			max_y = float(self._height)
			
		uniform = self._rnd.uniform
		
		if ball_type is None:
			balls = self._energy_balls
		else:
			balls = self._energy_balls_by_type[self._energy_balls_types[ball_type]]
			
		for ball in balls:
			ball.x = uniform(min_x, max_x)
			ball.y = uniform(min_y, max_y)
			
	def randomizeBallVelocities(self, x_range=None, y_range=None, ball_type=None):
		assert x_range is None or type(x_range) in (int, float)
		assert y_range is None or type(y_range) in (int, float)
		assert ball_type is None or isinstance(ball_type, basestring)
		
		if x_range is None:
			x_range = self._width / 2.0
		if y_range is None:
			y_range = self._height / 2.0
			
		uniform = self._rnd.uniform
		
		if ball_type is None:
			balls = self._energy_balls
		else:
			balls = self._energy_balls_by_type[self._energy_balls_types[ball_type]]
			
		for ball in balls:
			ball.vec_x = uniform(-x_range, x_range)
			ball.vec_y = uniform(-y_range, y_range)
			
	def tick(self):
		_rnd = self._rnd
		_width = float(self._width)
		_height = float(self._height)
		
		for (ball_type, balls) in self._energy_balls_by_type.iteritems():
			points = []
			for ball in balls:
				position = ball.tick(_width, _height, _rnd)
				if not position is None:
					points += position
			self._energy_balls_trails_by_type[ball_type] = [(GLfloat * len(points))(*points)] + self._energy_balls_trails_by_type[ball_type][:ball_type.trail_length]
			
	def draw(self):
		_ball_shader = self._ball_shader
		_trail_shader = self._trail_shader
		
		glMatrixMode(GL_PROJECTION)
		glPushMatrix()
		glLoadIdentity()
		gluOrtho2D(0.0, self._width, 0.0, self._height)
		glScalef(1.0, -1.0, 1.0)
		glTranslatef(0.0, -self._height, 0.0)
		glMatrixMode(GL_MODELVIEW)
		glLoadIdentity()
		
		glEnableClientState(GL_VERTEX_ARRAY)
		
		glEnable(GL_BLEND)
		glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
		_trail_shader.enable()
		for (ball_type, trails) in self._energy_balls_trails_by_type.iteritems():
			(r, g, b) = ball_type.colour
			_trail_shader.setUniform_vec3('rgb', r, g, b)
			
			trail_length = float(ball_type.trail_length)
			trail_brightness = ball_type.trail_brightness
			for (age, vertices) in enumerate(trails):
				_trail_shader.setUniform_float('a', (((1.0 - (float(age) / trail_length)) * trail_brightness) ** 2.0))
				glVertexPointer(2, GL_FLOAT, 0, vertices)
				glDrawArrays(GL_LINES, 0, len(vertices) // 2)
		_trail_shader.disable()
		glDisable(GL_BLEND)
		
		_ball_shader.enable()
		for (ball_type, balls) in self._energy_balls_by_type.iteritems():
			(r, g, b) = ball_type.colour
			_ball_shader.setUniform_vec3('rgb', r, g, b)
			glVertexPointer(2, GL_FLOAT, 0, ball_type.ball_vertices)
			
			for ball in balls:
				ball.draw()
		_ball_shader.disable()
		
		glDisableClientState(GL_VERTEX_ARRAY)
		
		glMatrixMode(GL_PROJECTION)
		glPopMatrix()
		glMatrixMode(GL_MODELVIEW)
		
	def group(self):
		for ball in self._energy_balls:
			ball.x = 640
			ball.y = 360
			ball.vec_x = ball.vec_y = 0.0
			ball.fixed_velocity = True
			
	def summon(self, x, y, ball_type):
		import random
		for ball in self._energy_balls_by_type[random.choice(self._energy_balls_types.values())]:
			ball.push((x - ball.x) / 7.0, (y - ball.y) / 7.0, True)
			
	def unsummon(self, ball_type):
		for ball in self._energy_balls:
			ball.release()
			
	def hide(self, ball_type):
		for ball in self._energy_balls_by_type[self._energy_balls_types[ball_type]]:
			ball.hide()
			
	def show(self, ball_type):
		for ball in self._energy_balls_by_type[self._energy_balls_types[ball_type]]:
			ball.show()
			
